package com.xu.rag.service.impl.model;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.json.JSONUtil;
import com.google.common.collect.Lists;
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.model.embeddings.EmbeddingRequest;
import com.volcengine.ark.runtime.service.ArkService;
import com.xu.rag.common.annotation.ModelPoint;
import com.xu.rag.common.annotation.ModelType;
import com.xu.rag.common.enums.ModelEnum;
import com.xu.rag.functionCall.GetWhetherTool;
import com.xu.rag.functionCall.SendMessageTool;
import com.xu.rag.pojo.dto.DoubaoToolDTO;
import com.xu.rag.pojo.result.ChunkResult;
import com.xu.rag.pojo.result.EmbeddingResult;
import com.xu.rag.service.LlmService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @Author: xdl
 * @CreateTime: 2025-01-13 23:22
 * @Description: 豆包
 */
@Slf4j
@Service
@ModelType(modelCode = "DOU_BAO")
public class DoubaoServiceImpl implements LlmService {

    private static final String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    private static final String embedding_point = "ep-20250205******-*****";
    private static final String chat_point = "ep-20250205******-*****";
    private static final String function_call_point = "ep-20250205******-*****";

    private static final String apiKey = "apiKey-******";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    static ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();

    @Override
    public String functionCall(String content) throws Exception {
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(content).build();
        messages.add(userMessage);
        final List<ChatTool> tools = Arrays.asList(
                new ChatTool(
                        "function",
                        new ChatFunction.Builder()
                        .name("get_current_whether")
                        .description("获取给定地点的天气")
                        .parameters(new DoubaoToolDTO(
                                "object",
                                new HashMap<String, Object>() {{
                                    put("location", new HashMap<String, String>() {{
                                        put("type", "string");
                                        put("description", "T地点的位置信息，比如北京");
                                    }});
                                    put("date", new HashMap<String, String>() {{
                                        put("type", "string");
                                        put("description", "日期或者时间，比如今天，比如2025-01-01");
                                    }});
                                }},
                                Lists.newArrayList("location", "date")
                        ))
                        .build()
                ),
                new ChatTool(
                        "function",
                        new ChatFunction.Builder()
                                .name("send_message")
                                .description("发送消息给指定的人")
                                .parameters(new DoubaoToolDTO(
                                        "object",
                                        new HashMap<String, Object>() {{
                                            put("name", new HashMap<String, String>() {{
                                                put("type", "string");
                                                put("description", "人名，比如小明");
                                            }});
                                        }},
                                        Lists.newArrayList("name")
                                ))
                                .build()
                )
        );

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(function_call_point)
                .messages(messages)
                .tools(tools)
                .build();

        // 查天气
        List<ChatCompletionChoice> whetherChoices = service.createChatCompletion(chatCompletionRequest).getChoices();
        for (ChatCompletionChoice choice : whetherChoices) {
            messages.add(choice.getMessage());
            for (ChatToolCall toolCall : choice.getMessage().getToolCalls()) {
                if ("get_current_whether".equals(toolCall.getFunction().getName())) {
                    GetWhetherTool getWhetherTool = JSONUtil.toBean(toolCall.getFunction().getArguments(), GetWhetherTool.class);
                    String call = getWhetherTool.call();

                    ChatMessage toolMessage = ChatMessage.builder()
                            .role(ChatMessageRole.TOOL)
                            .content(call)
                            .toolCallId(toolCall.getId())
                            .build();
                    messages.add(toolMessage);
                }
            }
        }

        // 发消息
        chatCompletionRequest.setMessages(messages);
        List<ChatCompletionChoice> sendMessageChoices = service.createChatCompletion(chatCompletionRequest).getChoices();
        for (ChatCompletionChoice choice : sendMessageChoices) {
            messages.add(choice.getMessage());
            for (ChatToolCall toolCall : choice.getMessage().getToolCalls()) {
                if ("send_message".equals(toolCall.getFunction().getName())) {
                    SendMessageTool sendMessageTool = JSONUtil.toBean(toolCall.getFunction().getArguments(), SendMessageTool.class);
                    String messageResult = sendMessageTool.sendMessage();
                    log.info("messageResult {}", messageResult);
                    ChatMessage toolMessage = ChatMessage.builder()
                            .role(ChatMessageRole.TOOL)
                            .content(messageResult)
                            .toolCallId(toolCall.getId())
                            .build();
                    messages.add(toolMessage);
                }
            }
        }
        // 总结
        chatCompletionRequest.setMessages(messages);
        List<ChatCompletionChoice> result = service.createChatCompletion(chatCompletionRequest).getChoices();
        return "最终答案：" + result.get(0).getMessage().getContent();
    }

    @Override
    public String chat(String prompt) {
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的AI人工智能助手").build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(prompt).build();
        messages.add(systemMessage);
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(chat_point)
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choices -> System.out.println(choices.getMessage().getContent()));
        return "";
    }

    @Override
//    @ModelPoint
    public double[] sentence(String sentence) {
        log.info("DOU_BAO [{}]", sentence);
        ChunkResult chunkResult = new ChunkResult();
        chunkResult.setChunkId(RandomUtil.randomInt());
        chunkResult.setContent(sentence);
        EmbeddingResult embeddingResult = this.embedding(chunkResult);
        return embeddingResult.getEmbedding();
    }

    @Override
    public List<EmbeddingResult> embedding(List<ChunkResult> chunkResults) {
        if (CollectionUtil.isEmpty(chunkResults)) {
            return new ArrayList<>();
        }
        List<EmbeddingResult> embeddingResults = new ArrayList<>();
        for (ChunkResult chunkResult : chunkResults) {
            embeddingResults.add(this.embedding(chunkResult));
        }
        return embeddingResults;
    }

    @Override
    public EmbeddingResult embedding(ChunkResult chunkResult) {
        List<String> inputs = new ArrayList<>();
        inputs.add(chunkResult.getContent());
        EmbeddingRequest chatCompletionRequest = EmbeddingRequest.builder()
                .model(embedding_point)
                .input(inputs)
                .build();

        com.volcengine.ark.runtime.model.embeddings.EmbeddingResult res = service.createEmbeddings(chatCompletionRequest);
        EmbeddingResult ret = new EmbeddingResult();
        List<Double> embedding = res.getData().get(0).getEmbedding();
        // 这里因为es dense vector范围为0-2048 暂且测试都1024
        List<Double> doubles = embedding.subList(0, 1024);
        double[] doubleArray = doubles.stream()
                .mapToDouble(Double::doubleValue)
                .toArray();
        ret.setEmbedding(doubleArray);
        ret.setPrompt(chunkResult.getContent());
        ret.setRequestId(Objects.toString(chunkResult.getChunkId()));
        return ret;
    }
}
